"""
动量策略是量化交易领域中极具实用价值的经典策略之一，其核心逻辑在于挖掘股票价格在特定周期内的延续性趋势。
该策略通过计算股票在过去 lookback_period（如默认 10 天）内的收益率，将其与预先设定的 buy_threshold（默认 10%）和 sell_threshold（默认 - 10%）进行对比，从而生成清晰的交易信号。
当收益率达到或超过买入阈值时，表明股价在短期内呈现强劲的上涨态势，策略即刻发出 “买入” 信号；
若收益率触及或低于卖出阈值，则意味着股价短期下跌趋势显著，此时策略给出 “卖出” 信号；
而当收益率介于两个阈值之间，股价处于相对平稳的波动状态，策略建议 “持有” 观望。
该策略通过 Python 代码实现，借助 SQLAlchemy 从数据库获取数据，
利用 matplotlib 以可视化表格呈现结果，并优化了布局与样式，使信号展示更为直观清晰。
策略逻辑简明高效，能有效捕捉股价短期趋势变化，适用于市场波动相对明显的行情环境，为投资者提供量化、客观的交易决策依据，通过灵活调整参数，还可适配不同股票的波动特性 。
"""
import pandas as pd
from sqlalchemy import create_engine
import matplotlib.pyplot as plt
from datetime import datetime
import numpy as np

# 创建数据库连接引擎
engine = create_engine('mysql+pymysql://root:@127.0.0.1:3306/stocks?charset=utf8mb4&use_unicode=1')

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False


def generate_momentum_signals(lookback_period=10, buy_threshold=0.10, sell_threshold=-0.10):
    """
    优化后的动量策略信号生成函数
    改进点：
    1. 缩小标题与表格间距
    2. 信号显示更紧凑（单行显示）
    3. 优化整体布局
    """
    # 获取股票数据
    query = "SELECT ts_code, name FROM stocks_info"
    stocks = pd.read_sql(query, con=engine)

    data = []
    for _, stock in stocks.iterrows():
        ts_code = stock['ts_code']
        name = stock['name']

        try:
            query = f"""
                SELECT trade_date, close 
                FROM stocks_daily 
                WHERE ts_code = '{ts_code}' 
                ORDER BY trade_date DESC 
                LIMIT {lookback_period + 5}
            """
            df = pd.read_sql(query, con=engine)

            if len(df) < lookback_period:
                signal = f"数据不足({lookback_period}天)"
            else:
                df = df.sort_values('trade_date')
                start_price = df['close'].iloc[0]
                end_price = df['close'].iloc[-1]
                return_pct = (end_price - start_price) / start_price * 100

                if return_pct >= buy_threshold * 100:
                    signal = f"买入(↑{return_pct:.2f}%)"
                elif return_pct <= sell_threshold * 100:
                    signal = f"卖出(↓{abs(return_pct):.2f}%)"
                else:
                    signal = f"持有({return_pct:+.2f}%)"

            data.append([name, ts_code, signal])

        except Exception as e:
            data.append([name, ts_code, f"错误:{str(e)}"])

    df_result = pd.DataFrame(data, columns=['股票名称', '股票代码', '信号'])

    # 动态计算图形尺寸
    num_rows = len(df_result)
    row_height = 0.5  # 缩小行高
    header_height = 0.8  # 缩小表头高度
    footer_space = 0.5  # 缩小底部留白
    title_space = 1.0  # 缩小标题区域高度

    fig_height = min(
        title_space + header_height + num_rows * row_height + footer_space,
        30  # 最大高度限制
    )

    # 创建图形
    fig = plt.figure(figsize=(12, fig_height))
    ax = fig.add_subplot(111)
    ax.axis('off')

    # 添加标题和时间戳（位置更靠近表格）
    current_time = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
    title = fig.suptitle(
        f'股票动量策略信号 ({lookback_period}日, 买入阈值:{buy_threshold * 100:.0f}%, 卖出阈值:{sell_threshold * 100:.0f}%)\n生成时间: {current_time}',
        y=0.98,  # 将标题位置下移更靠近表格
        fontsize=12,  # 缩小字体
        fontweight='bold'
    )

    # 计算表格位置（上移更靠近标题）
    table_height = 0.9 - (0.004 * num_rows)  # 调整表格高度

    # 创建表格
    table = ax.table(
        cellText=df_result.values,
        colLabels=df_result.columns,
        loc='center',
        bbox=[0.1, 0.05, 0.8, table_height],
        cellLoc='center'
    )

    # 设置表格样式
    table.auto_set_font_size(False)
    table.set_fontsize(10)

    # 设置行高（更紧凑）
    for i in range(len(df_result) + 1):
        for j in range(3):
            cell = table[(i, j)]
            cell.set_height(0.06)  # 缩小行高
            cell.set_edgecolor('lightgray')

    # 设置表头样式
    for j in range(3):
        table[(0, j)].set_facecolor('#40466e')
        table[(0, j)].get_text().set_color('white')
        table[(0, j)].set_fontsize(10)  # 缩小表头字体

    # 设置信号单元格颜色
    color_map = {
        '买入': '#ff6b6b',  # 红色
        '卖出': '#51cf66',  # 绿色
        '观望': '#339af0',  # 蓝色
    }

    for i in range(1, len(df_result) + 1):
        signal = df_result.iloc[i - 1]['信号']
        color = 'white'
        for key, val in color_map.items():
            if key in signal:
                color = val
                break

        for j in range(3):
            table[(i, j)].set_facecolor(color)
            if color != 'white':
                table[(i, j)].get_text().set_color('white')

    # 调整布局（减少上下边距）
    plt.tight_layout(rect=[0, 0, 1, 0.96])  # 上边距缩小到0.96
    plt.subplots_adjust(top=0.92, bottom=0.05)  # 顶部位置调整

    plt.show()


# 生成动量策略信号表格
generate_momentum_signals(lookback_period=10, buy_threshold=0.10, sell_threshold=-0.10)